package org.fatewa.engine.magic.db;

import lombok.SneakyThrows;
import org.apache.logging.log4j.util.Strings;
import org.fatewa.engine.magic.db.bean.Column;
import org.fatewa.engine.magic.db.bean.Database;
import org.fatewa.engine.magic.db.bean.Schema;
import org.fatewa.engine.magic.db.bean.Table;
import org.fatewa.engine.magic.db.enums.Dialect;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.stream.Collectors;

/**
 * 数据库扫描器，可进行数据表、数据表字段、数据库 Schema (如果有)
 * 扫描，支持 Mysql、Postgres、DM 8.x
 *
 * @author fatewa
 */
public class Scan {
    /**
     * 数据库方言
     */
    private final Dialect dialect;
    /**
     * 数据源
     */
    private final DataSource source;

    /**
     * 方言 MAP
     */
    private static final Map<Dialect, Map<String, String>> DIALECT_MAP = new HashMap<>();

    /**
     * 默认的方言
     */
    private static final HashMap<String, String> DEFAULT_DIALECT_MAP = new HashMap<>();

    /**
     * 默认数据库模式名称，无法查询到模式名称时使用
     */
    private static final String DEFAULT_SCHEMA_NAME = "default";

    /**
     * 类型为 表
     */
    private static final String TYPE_TABLE = "TABLE";
    /**
     * 表名称
     */
    private static final String TABLE_NAME = "TABLE_NAME";

    /**
     * 表所属 Schema
     */
    private static final String TABLE_SCHEM = "TABLE_SCHEM";
    /**
     * 备注（适用于表、字段等）
     */
    private static final String REMARK = "REMARKS";

    /**
     * 字段大小
     */
    private static final String COLUMN_SIZE = "COLUMN_SIZE";
    /**
     * 字段类型名称
     */
    private static final String COLUMN_TYPE_NAME = "TYPE_NAME";

    /**
     * 字段名称
     */
    private static final String COLUMN_NAME = "COLUMN_NAME";

    // 静态初始化方言
    static {
        DEFAULT_DIALECT_MAP.put(TYPE_TABLE, "TABLE");
        DEFAULT_DIALECT_MAP.put(TABLE_NAME, "TABLE_NAME");
        DEFAULT_DIALECT_MAP.put(TABLE_SCHEM, "TABLE_SCHEM");
        DEFAULT_DIALECT_MAP.put(REMARK, "REMARKS");
        // 设置默认方言
        DIALECT_MAP.put(Dialect.DEFAULT, DEFAULT_DIALECT_MAP);
    }

    /**
     * 创建扫描器之前需要指定 方言类型
     *
     * @param dialect 数据库方言
     * @param source  数据源配置
     */
    public Scan(Dialect dialect, DataSource source) {
        this.dialect = dialect;
        this.source = source;
    }

    /**
     * @param source 数据源配置
     */
    public Scan(DataSource source) {
        this(Dialect.DEFAULT, source);
    }

    /**
     * 扫描所有的表
     *
     * @return 所有表
     */
    public List<Schema> schemas() {
        return schemas("");
    }

    /**
     * 获取数据库元数据
     *
     * @return 数据库元数据
     */
    @SneakyThrows
    private DatabaseMetaData getMeta() {
        Connection connection = source.getConnection();
        return connection.getMetaData();
    }


    /**
     * 获取数据库方言
     *
     * @return 数据库方言
     */
    private Map<String, String> getDialect() {
        return DIALECT_MAP.get(dialect);
    }

    public Database database() {
        Database database = new Database();
        List<Schema> target;

        try {
            Map<String, String> currentDialect = getDialect();
            // 获取数据库元数据名称
            DatabaseMetaData meta = getMeta();

            String product = meta.getDatabaseProductName();
            database.setType(product);

            // 获取所有表
            ResultSet tables = meta.getTables(null, null, null, new String[]{currentDialect.get(TYPE_TABLE)});

            Map<String, Schema> schemaMap = new LinkedHashMap<>();
            List<Future<?>> futures = new ArrayList<>();
            ExecutorService service = Executors.newFixedThreadPool(10);

            while (tables.next()) {
                String name = tables.getString(currentDialect.get(TABLE_NAME));
                String schemaName = tables.getString(currentDialect.get(TABLE_SCHEM));
                if (Strings.isEmpty(schemaName)) {
                    // 设置默认名称
                    schemaName = DEFAULT_SCHEMA_NAME;
                }
                String remark = tables.getString(currentDialect.get(REMARK));

                // 初始化模式空间
                schemaMap.putIfAbsent(schemaName, new Schema(schemaName));
                if (Strings.isNotEmpty(name)) {
                    // 用于存储的 schema
                    Schema targetSchema = schemaMap.get(schemaName);
                    Table table = new Table();
                    table.setName(name).setRemark(remark);
                    targetSchema.table(table);
                    final String sn = schemaName;
                    Future<?> job = service.submit(() -> this.appendColumns(name, sn, table, meta));
                    futures.add(job);
                }
            }

            // 阻塞
            futures.forEach(it -> {
                try {
                    it.get();
                } catch (InterruptedException | ExecutionException e) {
                    throw new RuntimeException(e);
                }
            });

            target = new ArrayList<>(schemaMap.values());
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
        database.setSchemas(target);
        return database;
    }

    /**
     * 扫描一个模式下所有的表
     *
     * @param schema 数据库模式
     * @return 模式下的所有表
     */
    @SneakyThrows
    public List<Schema> schemas(String schema) {
        return database()
                .getSchemas()
                .stream().filter(it -> schema.equals(it.getName()))
                .collect(Collectors.toList());
    }

    /**
     * 将字段添加到表下
     *
     * @param tableName  表名称
     * @param schemaName 模式名称
     * @param table      表
     * @param meta       元数据
     */
    @SneakyThrows
    private void appendColumns(String tableName, String schemaName, Table table, DatabaseMetaData meta) {
        ResultSet columns = meta.getColumns(null, schemaName, tableName, null);
        // 遍历所有字段
        while (columns.next()) {
            String remark = columns.getString(REMARK);
            String name = columns.getString(COLUMN_NAME);
            String type = columns.getString(COLUMN_TYPE_NAME);
            String size = columns.getString(COLUMN_SIZE);
            Column column = new Column(name, remark, size, type);
            table.column(column);
        }
    }
}
